#include <iostream>
#include <fstream>
#include <string>
#include <cstring>
#include <chrono>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>

#define BUFFER_SIZE 1024
#define DEFAULT_PORT 8888

using namespace std;
using namespace std::chrono;

void send_file_tcp(const string& filename, const string& server_ip)
{
    int client_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (client_socket < 0)
	{
        perror("TCP socket creation failed");
        return;
    }

    sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(DEFAULT_PORT);
    inet_pton(AF_INET, server_ip.c_str(), &server_addr.sin_addr);

    if (connect(client_socket, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0)
	{
        perror("TCP connection failed");
        return;
    }

    ifstream infile(filename, ios::binary);
    if (!infile)
	{
        cerr << "Failed to open file: " << filename << endl;
        return;
    }

    char buffer[BUFFER_SIZE];
    auto start = high_resolution_clock::now();
    ssize_t total_bytes = 0;

    while (!infile.eof())
	{
        infile.read(buffer, BUFFER_SIZE);
        ssize_t bytes_read = infile.gcount();
        
        ssize_t bytes_sent = send(client_socket, buffer, bytes_read, 0);
        if (bytes_sent < 0)
		{
            perror("TCP send failed");
            break;
        }
        
        total_bytes += bytes_sent;
    }

    auto end = high_resolution_clock::now();
    auto duration = duration_cast<milliseconds>(end - start);
    double rate = (total_bytes / (1024.0 * 1024.0)) / (duration.count() / 1000.0);

    cout << "TCP Transfer completed. " << total_bytes << " bytes sent." << endl;
    cout << "Time: " << duration.count() << " ms" << endl;
    cout << "Transfer rate: " << rate << " MB/s" << endl;

    infile.close();
    close(client_socket);
}

void send_file_udp(const string& filename, const string& server_ip)
{
    int client_socket = socket(AF_INET, SOCK_DGRAM, 0);
    if (client_socket < 0)
	{
        perror("UDP socket creation failed");
        return;
    }

    sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(DEFAULT_PORT);
    inet_pton(AF_INET, server_ip.c_str(), &server_addr.sin_addr);

    ifstream infile(filename, ios::binary);
    if (!infile)
	{
        cerr << "Failed to open file: " << filename << endl;
        return;
    }

    char buffer[BUFFER_SIZE];
    auto start = high_resolution_clock::now();
    ssize_t total_bytes = 0;

    while (!infile.eof())
	{
        infile.read(buffer, BUFFER_SIZE);
        ssize_t bytes_read = infile.gcount();
        
        ssize_t bytes_sent = sendto(client_socket, buffer, bytes_read, 0, (struct sockaddr*)&server_addr, sizeof(server_addr));
        if (bytes_sent < 0)
		{
            perror("UDP send failed");
            break;
        }
        
        total_bytes += bytes_sent;
    }

    //发送信息
    const char* end_msg = "END";
    sendto(client_socket, end_msg, strlen(end_msg), 0, (struct sockaddr*)&server_addr, sizeof(server_addr));

    auto end = high_resolution_clock::now();
    auto duration = duration_cast<milliseconds>(end - start);
    double rate = (total_bytes / (1024.0 * 1024.0)) / (duration.count() / 1000.0);

    cout << "UDP Transfer completed. " << total_bytes << " bytes sent." << endl;
    cout << "Time: " << duration.count() << " ms" << endl;
    cout << "Transfer rate: " << rate << " MB/s" << endl;

    infile.close();
    close(client_socket);
}

int main(int argc, char* argv[])
{
    if (argc < 4)
	{
        cerr << "Usage: " << argv[0] << " <tcp|udp> <server_ip> <filename>" << endl;
        return 1;
    }

    string protocol = argv[1];
    string server_ip = argv[2];
    string filename = argv[3];

    if (protocol == "tcp") send_file_tcp(filename, server_ip);
    else if (protocol == "udp") send_file_udp(filename, server_ip);
    else
	{
        cerr << "Invalid protocol. Use 'tcp' or 'udp'." << endl;
        return 1;
    }

    return 0;
}